#ifndef BSPLINE_H_
#define BSPLINE_H_

#include <vector>
#include <cmath>
#include "polynomial.h"
#include "linearEqnSolver.h"
#include "Function.h"
using namespace std;

enum boundary_condition{
    order3,     // 针对3次complete边界条件
    order2,     // 针对2次边界节点处函数值相等
};

/**
 * 给定一组节点, 得到这组节点上的唯一样条基函数
 * e.g. 给出{-1,0,1,2,3},节点个数为5,唯一对应着一个3次样条基函数,紧支集为[-1,3].
 * 
 * 基函数由节点和存放着多项式的向量两个成员变量组成.给出节点时,在每段区间上
 * 存放一个多项式, 故用 vector 存放. 不难看出, knots 的元素个数比 polys 多 1.
 */
template <int order>
class basicBfun
{
private:
    vector<double> knots;              
    vector<polynomial<order>> polys;    
public:
    basicBfun(){};
    // 传入节点, 构造关于这组节点的样条插值基函数
    basicBfun(vector<double>& _knots);
    ~basicBfun(){};

    // 第i个节点与第j个节点的距离
    double ij_dist(int i, int j){ return abs(knots[i] - knots[j]); }

    // 打印样条基函数
    void print_BForm(ostream& out = cout);

    // 给定 x , 找出 x 所在区间
    int findIdx(double x);

    // 样条基函数在x处的值
    double get_value_at_x(double x);

    // 样条基函数在x处的导数值(基于讲义公式(3.41))
    double get_diff_at_x(double x);
};

/**
 * knots已知, B样条性质需要用到其他节点, 故定义延伸的节点extend_knots.
 * 延伸方式: 取给定节点knots所形成区间的平均区间长度, 等距延伸(与cardinal B相容)
 * 在每个子节点上构造样条基函数, 作为元素放入一个 vector 中.
 * 不难看出, extend_knots 的元素个数比 polys 多 order+1 个.
 */
template<int order>
class basicBfuns
{
private:
    vector<double> knots;               // e.g. 1 2 3 4 5 6 7 8, order = 3
    vector<double> extend_knots;        // -2 -1 0 1 2 3 4 5 6 7 8 9 10 11
    vector<basicBfun<order>> polys;    // 区间:[-2,2],[-1,3],...,[7,11]
public:
    basicBfuns(){};
    // 传入节点, 得到延伸节点和对应样条空间的一组基
    basicBfuns(const vector<double>& _knots);
    ~basicBfuns(){};

    // 打印给定节点的基
    void print_BForm(ostream& out = cout);

    // 打印采样点以作图
    void print_BForm_data(const vector<double>& coeff, 
                            int sample_num, ostream& out);

    // 给出文件名以存储密集散点
    void print_plot_data(const vector<double>& coeff, 
                            int sample_num, string &name);

    // 在每个基函数上x的值, 由于有 N+order-1 个多项式, 
    // 故返回一个长度为 N+order-1 的 vector
    vector<double> polys_value_at_x(double x);

    // 在每个基函数上x的导数值, 返回一个长度为 N+order-1 的 vector
    vector<double> polys_diff_at_x(double x);

    // 给定系数,在S(x)上x的值
    double value_with_coeff(double x, const vector<double>& coeff);
};

// 散点拟合求解
template <int order>
vector<double> solver(const vector<double>& knots, 
                        const vector<double>& sites, const vector<double>& vals,
                        const boundary_condition boundary, 
                        double left_bound, double right_bound)
{
    int N = knots.size();
    int unknows = N + order - 1;
    int _N = sites.size();
    if(unknows != _N+2 || sites.size() != vals.size()){
        cerr << "Error! 维数不匹配,请重试!" << endl;
        exit(-1);
    }

    basicBfuns<order> B(knots);
    // B.print_BForm();
    vector<double> b(vals);
    vector<vector<double>> A(_N, vector<double>(_N, 0.0));

    // 添加矩阵系数项(第一行和最后一行涉及边界条件,另行处理)
    for(int i = 1; i < _N-1; i++){
        vector<double> a_i = B.polys_value_at_x(sites[i]);
        a_i.erase(a_i.begin());
        a_i.erase(a_i.end()-1);
        A[i] = a_i;
    }

    if(boundary == order3){ // <-> order==3
        if(knots != sites){
            cout << "插值点不取在节点上时, 系数矩阵不为三对角矩阵, 这里不实现.\n";
            return vector<double>();
        }
    
        // 处理A的第一行
        vector<double> left = B.polys_value_at_x(sites[0]);
        vector<double> df_left = B.polys_diff_at_x(sites[0]);
        b[0] = df_left[0] * vals[0] - left[0] * left_bound;
        A[0][0] = df_left[0]*left[1] - left[0]*df_left[1];
        A[0][1] = df_left[0]*left[2] - left[0]*df_left[2];

        // 处理A的最后一行
        vector<double> right = B.polys_value_at_x(sites[_N-1]);
        vector<double> df_right = B.polys_diff_at_x(sites[_N-1]);
        b[_N-1] = df_right[unknows-1]*vals[_N-1] - right[unknows-1]*right_bound;
        A[_N-1][_N-1] = df_right[unknows-1]*right[unknows-2] 
                        - right[unknows-1]*df_right[unknows-2];
        A[_N-1][_N-2] = df_right[unknows-1]*right[unknows-3] 
                        - right[unknows-1]*df_right[unknows-3];

        vector<double> coeff = solu(A, b);  // 解方程, 还缺头尾两项
        double coeff_neg1 = (left_bound-coeff[0]*df_left[1]-coeff[1]*df_left[2]) 
                        / df_left[0];
        double coeff_N = (right_bound-coeff[_N-2]*df_right[unknows-3]
                    -coeff[_N-1]*df_right[unknows-2]) / df_right[unknows-1];
        coeff.push_back(coeff_N);
        coeff.insert(coeff.begin(), coeff_neg1);

        return coeff;

    }
    else if(boundary == order2){    // <-> order == 2
        vector<double> knot_left = B.polys_value_at_x(knots[0]);
        vector<double> site_left = B.polys_value_at_x(sites[0]);
        b[0] = vals[0]*knot_left[0] - site_left[0]*left_bound;
        A[0][0] = site_left[1]*knot_left[0]-knot_left[1]*site_left[0];
        A[0][1] = site_left[2] * knot_left[0];

        vector<double> knot_right = B.polys_value_at_x(knots[N-1]);
        vector<double> site_right = B.polys_value_at_x(sites[_N-1]);
        b[_N-1] = knot_right[unknows-1]*vals[_N-1]-site_right[unknows-1]*right_bound;
        A[_N-1][_N-1] = site_right[unknows-2]*knot_right[unknows-1] 
                        - site_right[unknows-1]*knot_right[unknows-2];
        A[_N-1][_N-2] = site_right[unknows-3] * knot_right[unknows-1];

        vector<double> coeff = solu(A, b);
        double coeff_0 = (left_bound-coeff[0]*knot_left[1]) / knot_left[0];
        double coeff_N = (right_bound-coeff[_N-1]*knot_right[unknows-2])
                         / knot_right[unknows-1];
        coeff.push_back(coeff_N);
        coeff.insert(coeff.begin(), coeff_0);

        return coeff;
    }
    else{
        cerr << "Error! 参数不匹配,请重试!\n";
        exit(-1);
    }
}

// 函数拟合求解
template <int order>
vector<double> solver(const vector<double>& knots, const vector<double>& sites,
                        const Function& fun, const boundary_condition boundary)
{
    vector<double> vals;
    int N = knots.size();
    int _N = sites.size();
    for(auto& site : sites){
        vals.push_back(fun.f(site));
    }
    double left_bound, right_bound;
    if(boundary == order3){
        left_bound = fun.df(sites[0]);
        right_bound = fun.df(sites[_N-1]);
    }
    else if(boundary == order2){
        left_bound = fun.f(knots[0]);
        right_bound = fun.f(knots[N-1]);
    }
    else{
        cerr << "Error! 参数不匹配,请重试!\n";
        exit(-1);
    }

    return solver<order>(knots, sites, vals, boundary, left_bound, right_bound);
}

/**
 * 包装.
 * 
 * 给上面的一组样条基函数basicBfuns配上系数, 成为某个函数的拟合表达式(或某组散点的拟合)
 */ 
template <int order>
class BSpline
{
private:
    vector<double> coeffs;
    basicBfuns<order> Bfuns;
public:
    BSpline(){};
    BSpline(const vector<double>& knots, 
            const vector<double>& sites, const vector<double>& vals,
            const boundary_condition boundary, 
            double left_bound, double right_bound)
        : coeffs(solver<order>(knots, sites, vals, boundary, 
                left_bound, right_bound)),
          Bfuns(basicBfuns<order>(knots)) {};
    BSpline(const vector<double>& knots, const vector<double>& sites,
            const Function& fun, const boundary_condition boundary)
        : coeffs(solver<order>(knots, sites, fun, boundary)),
          Bfuns(basicBfuns<order>(knots)) {};
    ~BSpline(){};

    // B样条在x处的值
    double get_value_at_x(double x){
        return Bfuns.value_with_coeff(x, coeffs);
    }

    // 给出文件名和样本点数, 以存储密集散点, 方便作图
    void print_plot_data(int sample_num, string &name){
        Bfuns.print_plot_data(coeffs, sample_num, name);
    }

};


// 函数具体实现

template <int order>
basicBfun<order>::basicBfun(vector<double>& _knots): knots(_knots){
    int N = knots.size();
    if(N - 2 != order){
        cerr << "Error! 维数不匹配!\n";
        exit(-1);
    }
    // 这里暴力穷举了, 没采用递归, 且只穷举到3维情况.
    if(order == 0){
        polys.push_back( polynomial<order>({1}) );
    }
    else if(order == 1){
        polynomial<order> p1({ -knots[0] / ij_dist(0,1),
                                1.0 / ij_dist(0,1) });
        polynomial<order> p2({ knots[2] / ij_dist(1,2),
                                -1.0 / ij_dist(1,2) });
        polys.push_back(p1);
        polys.push_back(p2);
    }
    else if(order == 2){
        double divisor = ij_dist(0,1) * ij_dist(0,2);
        double co0 = knots[0]*knots[0] / divisor;
        double co1 = -2.0*knots[0] / divisor;
        double co2 = 1.0 / divisor;
        polys.push_back( polynomial<order>({co0, co1, co2}) );
    
        double divisor1 = ij_dist(0,2) * ij_dist(1,2);
        double divisor2 = ij_dist(1,2) * ij_dist(1,3);
        co0 = -knots[2]*knots[0]/divisor1
            - knots[3]*knots[1]/divisor2;
        co1 = (knots[2]+knots[0])/divisor1
            + (knots[3]+knots[1])/divisor2;
        co2 = -1.0/divisor1 - 1.0/divisor2;
        polys.push_back( polynomial<order>({co0, co1, co2}) );
    
        divisor = (knots[3]-knots[1])*(knots[3]-knots[2]);
        co0 = knots[3]*knots[3] / divisor;
        co1 = -2.0*knots[3] / divisor;
        co2 = 1.0 / divisor;
        polys.push_back( polynomial<order>({co0, co1, co2}) );
    }
    else if(order == 3){
        vector<double> coeff(N);
        double divisor = ij_dist(0,1) * ij_dist(0,2) * ij_dist(0,3);
        coeff[0] = - pow(knots[0], 3) / divisor;
        coeff[1] = 3.0 * pow(knots[0], 2) / divisor;
        coeff[2] = -3.0 * knots[0] / divisor;
        coeff[3] = 1.0 / divisor;
        polys.push_back( polynomial<order>(coeff));

        double divisor1 = ij_dist(0,3) * ij_dist(0,2) * ij_dist(1,2);
        double divisor2 = ij_dist(0,3) * ij_dist(1,3) * ij_dist(1,2);
        double divisor3 = ij_dist(1,4) * ij_dist(1,3) * ij_dist(1,2);
        coeff[0] = pow(knots[0],2)*knots[2] / divisor1
                + knots[0]*knots[1]*knots[3] / divisor2
                + knots[4]*pow(knots[1],2) / divisor3;
        coeff[1] = (-pow(knots[0],2)-2.0*knots[0]*knots[2]) / divisor1
                - (knots[0]*knots[1]+knots[0]*knots[3]+knots[1]*knots[3])/divisor2
                + (-pow(knots[1],2)-2.0*knots[1]*knots[4]) / divisor3;
        coeff[2] = (knots[2]+2.0*knots[0]) / divisor1
                + (knots[0]+knots[1]+knots[3]) / divisor2
                + (2.0*knots[1]+knots[4]) / divisor3;
        coeff[3] = -1.0/divisor1 - 1.0/divisor2 - 1.0/divisor3;
        polys.push_back( polynomial<order>(coeff));

        divisor1 = ij_dist(0,3) * ij_dist(1,3) * ij_dist(2,3);
        divisor2 = ij_dist(1,4) * ij_dist(1,3) * ij_dist(2,3);
        divisor3 = ij_dist(1,4) * ij_dist(2,4) * ij_dist(2,3);
        coeff[0] = -pow(knots[3],2)*knots[0] / divisor1
                - knots[1]*knots[3]*knots[4] / divisor2
                - knots[2]*pow(knots[4],2) / divisor3;
        coeff[1] = (pow(knots[3],2)+2.0*knots[3]*knots[0]) / divisor1
                + (knots[3]*knots[4]+knots[1]*knots[3]+knots[1]*knots[4])/divisor2
                + (pow(knots[4],2)+2.0*knots[4]*knots[2]) / divisor3;
        coeff[2] = (-knots[0]-2.0*knots[3]) / divisor1
                - (knots[1]+knots[3]+knots[4]) /divisor2
                - (2.0*knots[4]+knots[2]) / divisor3;
        coeff[3] = 1.0/divisor1 + 1.0/divisor2 + 1.0/divisor3;
        polys.push_back( polynomial<order>(coeff));

        divisor = ij_dist(1,4) * ij_dist(2,4) * ij_dist(3,4);
        coeff[0] = pow(knots[4], 3) / divisor;
        coeff[1] = -3.0 * pow(knots[4], 2) / divisor;
        coeff[2] = 3.0 * knots[4] / divisor;
        coeff[3] = -1.0 / divisor;
        polys.push_back( polynomial<order>(coeff));
    }
    else{
        cerr << "样条基函数暂不支持" << order << "阶及以上次数!\n";
        exit(-1);
    }
}

template <int order>
void basicBfun<order>::print_BForm(ostream& out){
    int N = knots.size();
    for(int i = 0; i < N-1; i++){
        out << "区间[" << knots[i] << "," << knots[i+1] << "]:";
        polys[i].print_poly();
    }
}

template <int order>
int basicBfun<order>::findIdx(double x){
    int N = knots.size();
    for(int i = 0; i < N-1; i++){
        if(knots[i] <= x && x <= knots[i+1])
            return i;
    }
    return -1;  // 不在节点给定区间内, 返回-1.
}

template <int order>
double basicBfun<order>::get_value_at_x(double x){
    int index = findIdx(x);
    if(index == -1) return 0;
    return polys[index].get_value_at_x(x);
}

template <int order>
double basicBfun<order>::get_diff_at_x(double x){
    int N = knots.size();
    vector<double> left_knots, right_knots;
    left_knots.assign(knots.begin(), knots.end()-1);
    right_knots.assign(knots.begin()+1, knots.end());
    basicBfun<order-1> left_basic(left_knots);
    basicBfun<order-1> right_basic(right_knots);
    double left = order * left_basic.get_value_at_x(x) / ij_dist(0, N-2);
    double right = order * right_basic.get_value_at_x(x) / ij_dist(1, N-1);
    return left - right;
} 

template <int order>
basicBfuns<order>::basicBfuns(const vector<double>& _knots): knots(_knots)
{
    int N = knots.size();
    extend_knots = _knots;
    double average = (knots[N-1]-knots[0]) / (N-1); // 取平均区间长度
    double head = extend_knots[0];
    double tail = extend_knots[N-1];
    // 开始延伸
    for(int i = 0; i < order; i++){
        head -= average;
        extend_knots.insert(extend_knots.begin(), head);
        tail += average;
        extend_knots.push_back(tail);
    }

    for(int first_knot = 0; first_knot < N+order-1; first_knot++){
        vector<double> subknots;    // 求一段子节点上对应的基函数
        for(int i = 0; i < order+2; i++){
            subknots.push_back(extend_knots[first_knot + i]);
        }
        polys.push_back(basicBfun<order>(subknots));
    }
}

template <int order>
void basicBfuns<order>::print_BForm(ostream& out){
    int N = knots.size();
    for(int i = 0; i < N+order-1; i++){
        out << "区间[" << extend_knots[i] << "," 
            << extend_knots[i+order+1] << "]上的样条基函数:" << endl;
        polys[i].print_BForm(out);
    }
} 

template <int order>
void basicBfuns<order>::print_BForm_data(const vector<double>& coeff, 
                            int sample_num, ostream& out)
{
    if(coeff.size() != polys.size()){
        cerr << "Error! 参数不匹配!\n";
        exit(-1);
    }
    int N = knots.size();
    double delta = (knots[N-1] - knots[0]) / sample_num;
    double x = knots[0];
    for(int i = 0; i <= sample_num; i++){
        out << x << " " << value_with_coeff(x, coeff) << endl;
        x += delta; 
    }
}

template <int order>
void basicBfuns<order>::print_plot_data(const vector<double>& coeff, 
                            int sample_num, string &name)
{
    ofstream file(name);
    if(file.is_open()){
        print_BForm_data(coeff, sample_num, file);
        file.close();
        cout << "File " << name;
        cout << " created and written successfully.\n";
    }
    else{  
        cout << "Error creating file: " << name << ".\n";  
    }
}


template <int order>
vector<double> basicBfuns<order>::polys_value_at_x(double x){
    vector<double> values;
    for(auto& poly : polys){
        values.push_back(poly.get_value_at_x(x));
    }
    return values;
}

template <int order>
vector<double> basicBfuns<order>::polys_diff_at_x(double x){
    vector<double> diffs;
    for(auto& poly : polys){
        diffs.push_back(poly.get_diff_at_x(x));
    }
    return diffs;
}

template <int order>
double basicBfuns<order>::value_with_coeff(double x, 
                                            const vector<double>& coeff){
    if(coeff.size() != polys.size()){
        cerr << "Error! 参数不匹配!\n";
        exit(-1);
    }
    int N = coeff.size();
    double value = 0.0;
    for(int i = 0; i < N; i++){
        value += coeff[i] * polys[i].get_value_at_x(x);
    }
    return value;
}

#endif